package com.itmatter.classifier.directdb;

import java.util.Collections;
import java.util.Set;

import com.enigmastation.extractors.WordLister;
import com.enigmastation.extractors.impl.StemmingWordLister;
import com.itmatter.classifier.DBClassifier;
import com.itmatter.classifier.util.ClassifierHibernateUtil;
import com.itmatter.classifier.domain.ClassifierCategory;
import com.itmatter.classifier.domain.ClassifierFeature;
import com.itmatter.classifier.domain.ClassifierFeatureId;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import org.apache.log4j.Logger;
import org.hibernate.Criteria;
import org.hibernate.FlushMode;
import org.hibernate.HibernateException;
import org.hibernate.Query;
import org.hibernate.SQLQuery;
import org.hibernate.Session;
import org.hibernate.Transaction;

/**
 * Based on initial work from ci-bayes-2.1.0-SNAPSHOT-sources
 * @author Daniel Fisla, dfisla@itmatter.com, ITMATTER Inc., http://www.itmatter.com, Copyright 2010, LGPLv3
 */
public class DBClassifierImpl implements DBClassifier {

    private static final Logger logger = Logger.getLogger(DBClassifierImpl.class.getName());
    protected WordLister wordLister = null;
    protected boolean initialized = false;
    private double WEIGHT = 1.0;
    private double ASSUMED_PROBABILITY = 0.5;

    /**
     * 
     */
    public DBClassifierImpl() {
        init();
    }

    /**
     * 
     */
    public synchronized void init() {
        if (!initialized) {

            if (wordLister == null) {
                wordLister = new StemmingWordLister();
            }

            initialized = true;
        }
    }

    /**
     * 
     * @return
     */
    public WordLister getWordLister() {
        return wordLister;
    }

    /**
     * 
     * @param wordLister
     */
    public void setWordLister(WordLister wordLister) {
        this.wordLister = wordLister;
    }

    /**
     * Increase the count of a feature.
     * 
     * @param feature
     *            the feature (the 'word')
     * @param category
     *            the category
     */
    public void incrementFeature(String feature, String category) {

        Transaction tx = null;
        try {
            Session session = ClassifierHibernateUtil.getSessionFactory().getCurrentSession();
            tx = session.beginTransaction();
            session.setFlushMode(FlushMode.COMMIT);

            ClassifierFeature cFeature = (ClassifierFeature) session.get(ClassifierFeature.class, new ClassifierFeatureId(feature, category));

            if (cFeature == null) {
                cFeature = new ClassifierFeature(new ClassifierFeatureId(feature, category));
            }

            cFeature.setCounts(cFeature.getCounts() + 1);
            session.saveOrUpdate(cFeature);
            tx.commit();

        } catch (HibernateException eh) {
            try {
                if (tx != null) {
                    tx.rollback();
                    logger.error("ROLLBACK!");
                }
            } catch (HibernateException eh2) {
            }

            logger.error(eh.getMessage(), eh);
        }

    }

    /**
     * Increase the count of a category. 
     * 
     * @param category
     *            the category to increment
     */
    public void incrementCategory(String category) {

        ClassifierCategory cCategory = null;
        Transaction tx = null;
        try {
            Session session = ClassifierHibernateUtil.getSessionFactory().getCurrentSession();
            tx = session.beginTransaction();
            session.setFlushMode(FlushMode.COMMIT);

            cCategory = (ClassifierCategory) session.get(ClassifierCategory.class, category);

            if (cCategory == null) {
                cCategory = new ClassifierCategory(category);
            }
            cCategory.setCounts(cCategory.getCounts() + 1);
            session.saveOrUpdate(cCategory);
            tx.commit();
        } catch (HibernateException eh) {
            try {
                if (tx != null) {
                    tx.rollback();
                    logger.error("ROLLBACK!");
                }
            } catch (HibernateException eh2) {
            }

            logger.error(eh.getMessage(), eh);
        }


    }

    /**
     * 
     * @param feature
     *            the feature
     * @param category
     *            the category to query
     * @return the number of times a feature has appeared in a category
     */
    public double getFeatureCount(String feature, String category) {

        double rv = 0.0;

        Transaction tx = null;
        try {
            Session session = ClassifierHibernateUtil.getSessionFactory().getCurrentSession();
            tx = session.beginTransaction();
            session.setFlushMode(FlushMode.COMMIT);

            ClassifierFeature cFeature = (ClassifierFeature) session.get(ClassifierFeature.class, new ClassifierFeatureId(feature, category));

            if (cFeature != null) {
                rv = cFeature.getCounts();
            }
            tx.commit();

        } catch (HibernateException eh) {
            try {
                if (tx != null) {
                    tx.rollback();
                    logger.error("ROLLBACK!");
                }
            } catch (HibernateException eh2) {
            }

            logger.error(eh.getMessage(), eh);
        }

        return rv;
    }

    /**
     * 
     * @param category
     *            the category to count items for
     * @return the number of items in a category
     */
    public double getCategoryCount(String category) {

        double rv = 0.0;
        Transaction tx = null;
        try {
            Session session = ClassifierHibernateUtil.getSessionFactory().getCurrentSession();
            tx = session.beginTransaction();
            session.setFlushMode(FlushMode.COMMIT);

            ClassifierCategory cCategory = (ClassifierCategory) session.get(ClassifierCategory.class, category);

            if (cCategory != null) {
                rv = (double) cCategory.getCounts();
            }
            tx.commit();

        } catch (HibernateException eh) {
            try {
                if (tx != null) {
                    tx.rollback();
                    logger.error("ROLLBACK!");
                }
            } catch (HibernateException eh2) {
            }

            logger.error(eh.getMessage(), eh);
        }

        return rv;
    }

    /**
     * 
     * 
     * @return the total number of items
     */
    public double getTotalCount() {

        double totalCategoryCount = 0.0;
        Transaction tx = null;
        try {
            Session session = ClassifierHibernateUtil.getSessionFactory().getCurrentSession();
            tx = session.beginTransaction();
            session.setFlushMode(FlushMode.COMMIT);

            Criteria criteria = session.createCriteria(ClassifierCategory.class);
            List rv = criteria.list();

            if (rv != null) {
                for (Iterator it = rv.iterator(); it.hasNext();) {
                    ClassifierCategory cCategory = (ClassifierCategory) it.next();
                    totalCategoryCount += cCategory.getCounts();
                }
            }
            tx.commit();

        } catch (HibernateException eh) {
            try {
                if (tx != null) {
                    tx.rollback();
                    logger.error("ROLLBACK!");
                }
            } catch (HibernateException eh2) {
            }

            logger.error(eh.getMessage(), eh);
        }

        return totalCategoryCount;
    }

    /**
     * 
     * 
     * @param feature
     * @return
     */
    public double getTotalFeatureCount(String feature) {

        double totalFeatureCount = 0.0;
        Transaction tx = null;
        try {
            Session session = ClassifierHibernateUtil.getSessionFactory().getCurrentSession();
            tx = session.beginTransaction();
            session.setFlushMode(FlushMode.COMMIT);

            SQLQuery q = session.createSQLQuery("SELECT * FROM cls_features WHERE feature = '" + feature + "'").addEntity(ClassifierFeature.class);
            List rv = q.list();

            if (rv != null) {        
                for (Iterator it = rv.iterator(); it.hasNext();) {
                    ClassifierFeature cFeature = (ClassifierFeature) it.next();
                    totalFeatureCount += cFeature.getCounts();
                }               
            }
            tx.commit();

        } catch (HibernateException eh) {
            try {
                if (tx != null) {
                    tx.rollback();
                    logger.error("ROLLBACK!");
                }
            } catch (HibernateException eh2) {
            }

            logger.error(eh.getMessage(), eh);
        }

        return totalFeatureCount;
    }

    /**
     * 
     * 
     * @return the list of all getCategories
     */
    public Set<String> getCategories() {

        Set<String> cats = new HashSet<String>();
        Transaction tx = null;
        try {
            Session session = ClassifierHibernateUtil.getSessionFactory().getCurrentSession();
            tx = session.beginTransaction();
            session.setFlushMode(FlushMode.COMMIT);

            Criteria criteria = session.createCriteria(ClassifierCategory.class);
            List rv = criteria.list();

            if (rv != null) {
                for (Iterator it = rv.iterator(); it.hasNext();) {
                    ClassifierCategory cat = (ClassifierCategory) it.next();
                    cats.add(cat.getCategory());
                }
            }
            tx.commit();

        } catch (HibernateException eh) {
            try {
                if (tx != null) {
                    tx.rollback();
                    logger.error("ROLLBACK!");
                }
            } catch (HibernateException eh2) {
            }

            logger.error(eh.getMessage(), eh);
        }

        return Collections.unmodifiableSet(cats);
    }

    /**
     * 
     * @param item
     * @param category
     */
    public void train(Object item, String category) {
        Set<String> features = wordLister.getUniqueWords(item);

        for (String f : features) {
            incrementFeature(f, category);
        }
        incrementCategory(category);
    }

    /**
     * @param feature
     *            the feature to consider
     * @param category
     *            the category
     * @return the feature probability for the class
     */
    public double getFeatureProbability(String feature, String category) {
        try {
            return getFeatureCount(feature, category) / getCategoryCount(category);
        } catch (ArithmeticException ae) {
            return 0;
        }
    }

    /**
     * @param feature
     *            The feature to consider
     * @param category
     *            the category to consider weight for
     * @return the weighted probability
     */
    public double getWeightedFeatureProbability(String feature, String category) {
        double basicprob = getFeatureProbability(feature, category);
        //System.out.println("getFeatureProbability FEA: " + feature + " basicprob: " + basicprob);
        double totals = getTotalFeatureCount(feature);
        //System.out.println("getTotalFeatureCount FEA: " + feature + " totals: " + totals);
        return ((WEIGHT * ASSUMED_PROBABILITY) + (totals * basicprob)) / (WEIGHT + totals);
    }

    /**
     * 
     * 
     */
    public void deleteAllData() {
        logger.info("DELETE-ALL-DATA!");
        Transaction tx = null;
        try {
            Session session = ClassifierHibernateUtil.getSessionFactory().getCurrentSession();
            tx = session.beginTransaction();
            session.setFlushMode(FlushMode.COMMIT);

            String hqlCategories = String.format("delete from ClassifierCategory");
            Query queryCategories = session.createQuery(hqlCategories);

            int rv = queryCategories.executeUpdate();
            logger.info("DELETED CLS_CATEGORIES records: " + rv);

            tx.commit();
        } catch (HibernateException eh) {
            try {
                if (tx != null) {
                    tx.rollback();
                    logger.error("ROLLBACK!");
                }
            } catch (HibernateException eh2) {
            }

            logger.error(eh.getMessage(), eh);
        } finally {
        }

        try {
            Session session = ClassifierHibernateUtil.getSessionFactory().getCurrentSession();
            tx = session.beginTransaction();
            session.setFlushMode(FlushMode.COMMIT);

            String hqlFeatures = String.format("delete from ClassifierFeature");
            Query queryFeatures = session.createQuery(hqlFeatures);

            int rv = queryFeatures.executeUpdate();
            logger.info("DELETED CLS_FEATURES records: " + rv);

            tx.commit();
        } catch (HibernateException eh) {
            try {
                if (tx != null) {
                    tx.rollback();
                    logger.error("ROLLBACK!");
                }
            } catch (HibernateException eh2) {
            }

            logger.error(eh.getMessage(), eh);
        } finally {
        }
    }
}
